from functools import wraps
from typing import Any, Sequence, Union

from flask import current_app, request
from flask_jwt_extended import get_jwt_identity, verify_jwt_in_request

from app.models import UserInfo
from app.utils.exception import BaseError


LocationType = Union[str, Sequence, None]


def jwt_required(
        optional: bool = False,
        fresh: bool = False,
        refresh: bool = False,
        locations: LocationType = None,
) -> Any:
    def wrapper(fn):
        @wraps(fn)
        def decorator(*args, **kwargs):
            try:
                verify_jwt_in_request(
                    optional, fresh, refresh, locations
                )
            except Exception as e:
                print(e)
                return BaseError(msg='登录超时').jsonify()

            user_info = get_jwt_identity()
            user_id = user_info.get('user_id')

            # To check if a user is logged in
            if not UserInfo.queryset().filter(UserInfo.id == user_id).count() > 0:
                return BaseError(msg='当前用户不存在').jsonify()

            return current_app.ensure_sync(fn)(*args, **kwargs)

        return decorator

    return wrapper
